# encoding: utf-8

import numpy as np
import torch
from torch.utils import data


# dataloader
class DataGen(data.Dataset):
    def __init__(self, data, label):
        self.data = data
        self.label = label

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return np.array(self.data[index]), np.array(self.label[index])


class ModelExpert(torch.nn.Module):
    def __init__(self, bert_model, bert_config, num_class: int, dropout: float = 0.1):
        super(ModelExpert, self).__init__()

        self.bert_model = bert_model
        self.dropout = torch.nn.Dropout(dropout)
        self.fc = torch.nn.Linear(bert_config.hidden_size, num_class)

        self.leaky_relu = torch.nn.LeakyReLU()

    def forward(self, token_ids):
        bert_out = self.bert_model(token_ids)[1]
        bert_out = self.dropout(bert_out)
        bert_out = self.leaky_relu(bert_out)
        bert_out = self.fc(bert_out)

        return bert_out


class Gating(torch.nn.Module):
    def __init__(self, input_dim, num_experts):
        super(Gating, self).__init__()

        # Layers
        self.layer1 = torch.nn.Linear(input_dim, num_experts)

    def forward(self, x):
        x = self.layer1(x)
        x = torch.softmax(x, dim=1)
        return x


class ModelMoe(torch.nn.Module):
    def __init__(self, trained_experts):
        super(ModelMoe, self).__init__()
        self.experts = torch.nn.ModuleList(trained_experts)
        num_experts = len(trained_experts)
        # Assuming all experts have the same input dimension
        input_dim = trained_experts[0].bert_model.in_features
        self.gating = Gating(input_dim, num_experts)

    def forward(self, x):
        # Get the weights from the gating network
        weights = self.gating(x)

        # Calculate the expert outputs
        outputs = torch.stack(
            [expert(x) for expert in self.experts], dim=2)

        # Adjust the weights tensor shape to match the expert outputs
        weights = weights.unsqueeze(1).expand_as(outputs)

        # Multiply the expert outputs with the weights and
        # sum along the third dimension
        return torch.sum(outputs * weights, dim=2)

